/*
 *        (C) COPYRIGHT Ingenic Limited.
 *             ALL RIGHTS RESERVED
 *
 * File       : inference_yolov8s.cpp
 * Authors    : ljhong, cywang
 * Create Time: 2025-01-23:15:57:21
 * Description:
 * 
 */

 #define STB_IMAGE_IMPLEMENTATION
 #include "./stb/stb_image.h"
 #include "./stb/drawing.hpp"
 #define STB_IMAGE_WRITE_IMPLEMENTATION
 #include "./stb/stb_image_write.h"
 static const uint8_t color[3] = {0xff, 0, 0};
 
 #include "ingenic_mnni.h"
 #include "common/common_utils.h"
 #include "utils/imgproc.h"
 #include "utils/postproc.h"
 
 #include <cmath>
 #include <errno.h>
 #include <fstream>
 #include <iostream>
 #include <memory>
 #include <stdio.h>
 #include <stdlib.h>
 #include <string>
 #include <cstring>
 #include <string.h>
 #include <vector>
 #include <algorithm>
 #include <dirent.h>
 #include <sys/stat.h>
 #include <unistd.h>
 
 #define IS_ALIGN_64(x) (((size_t)x) & 0x3F)
 
 int input_w = 640;
 int input_h = 640;
 int ori_img_w;
 int ori_img_h;
 float person_threshold = 0.5;
 float nms_threshold = 0.6;
 int classes = 1;
 std::vector<int> strides = {8,16,32};
 int reg_max = 16;
 
 using namespace std;
 using namespace magik::uranus;
 
 #ifdef URANUS_NV12
 magik::uranus::ChannelLayout DATA_FORMAT = magik::uranus::ChannelLayout::NV12;
 #else
 magik::uranus::ChannelLayout DATA_FORMAT = magik::uranus::ChannelLayout::RGBA;
 #endif
 
 typedef unsigned char uint8_t;
 #define RET_CHECKER(ret, msg)                                           \
     do {                                                                \
         if (!(ret)) {                                                   \
             fprintf(stderr, "%s [%s %d]\n", msg, __func__, __LINE__);   \
             return -1;                                                  \
         }                                                               \
     } while(0)
 
 typedef struct {
     uint8_t *data;
     int width;
     int height;
 }FrameInfo;
 
 static int fill_dim(uint8_t *data, int h, int w) {
     for (int i = 0; i < h*w; i++) {
         data[i*4+3] = 0;
     }
     return 0;
 }
 
 int rgb2bgr(uint8_t *data, int data_cnt){
     uint8_t* p = data;
     int pixel_c = data_cnt / 4;
     for (int i=0; i < pixel_c; ++i) {
         uint8_t t = p[0];
         p[0] = p[2];
         p[2] = t;
         p += 4;
     }
 }
 
 void read_img(FrameInfo &in, Img &img, const std::string &img_path) {
     int comp = 0;
     int w, h, c;
     unsigned char *imagedata = stbi_load(img_path.c_str(), &ori_img_w, &ori_img_h, &comp, 3); // image format is rgba
     in.data = stbi_load(img_path.c_str(), &w, &h, &c, 4);
     in.width = w;
     in.height = h;
     printf("input===> %d %d \n",in.width, in.height);
     int _size = in.width * in.height * 4;
     // printf("rgb ==> first channel => %d %d %d \n", in.data[0], in.data[1], in.data[2]);
     // rgb2bgr(in.data, _size);
     // printf("rgb ==> bgr first channel => %d %d %d \n", in.data[0], in.data[1], in.data[2]);
     img.w = ori_img_w;
     img.h = ori_img_h;
     img.c = 3;
     img.w_stride = ori_img_w * 3;
     img.data = imagedata;
     fill_dim(in.data, w, h);
 }
 
 
 void prepare_data(FrameInfo &in, magik::uranus::Tensor &out) {
     int in_w = in.width;
     int in_h = in.height;
     printf("perpare_data, in_h in_w ==>%d %d", in.height, in.width);
     magik::uranus::shape_t in_shape;
     magik::uranus::DataFormat data_format;
     if (DATA_FORMAT == magik::uranus::ChannelLayout::NV12) {
         in_shape = {1, in_h, in_w, 1};
         data_format = magik::uranus::DataFormat::NV12;
         
     } else {
         in_shape = {1, in_h, in_w, 4};
         data_format = magik::uranus::DataFormat::NHWC;
     }
     auto in_tensor = magik::uranus::Tensor(in_shape, magik::uranus::DataType::UINT8, data_format);
     uint8_t *in_tensor_data = in_tensor.mudata<uint8_t>();
     memcpy((void *)in_tensor_data, (void *)in.data, in_tensor.get_bytes_size());
 
     // resize
     magik::uranus::BsCommonParam param;
     param.pad_val = 114;
     param.pad_type = magik::uranus::BsPadType::SYMMETRY;
     param.in_layout = DATA_FORMAT;
     param.out_layout = magik::uranus::ChannelLayout::RGBA;
     param.input_height = in_h;
     param.input_width = in_w;
     param.input_line_stride = (DATA_FORMAT == magik::uranus::ChannelLayout::NV12) ? in_w : in_w * 4;
     param.addr_attr.vir_addr = in_tensor_data;
     vector<magik::uranus::Box_t> output_tmp;
     magik::uranus::Box_t tmp_box{0, 0, in_w, in_h};
     output_tmp.push_back(tmp_box);
     printf("work_h, %d work_w:%d \n", input_h, input_w);
     out.reshape({1, input_h, input_w, 4});
     uint8_t *out_data = out.mudata<uint8_t>();
     vector<magik::uranus::Tensor> output_tesnor;
     output_tesnor.push_back(out);
     magik::uranus::crop_common_resize(output_tesnor, output_tmp, magik::uranus::AddressLocate::NMEM_VIRTUAL, &param);
 
     if (IS_ALIGN_64(out_data) != 0) {
         fprintf(stderr, "input addr not align to 64 bytes.\n");
         return ;
     }
 }
 
 void generateBBox(std::vector<magik::uranus::Tensor> out_res, std::vector<magik::uranus::ObjBbox_t>& candidate_boxes) {   
     magik::uranus::generateBBox_yolov8(out_res, candidate_boxes, strides, person_threshold, classes, reg_max);
 }

 
 void scale_boxes(int ori_w, int ori_h, int new_w, int new_h, std::vector<magik::uranus::ObjBbox_t> &out_boxes) {
     float scale;
     float pad_x, pad_y;
     scale = std::min((float(new_w) / ori_w), (float(new_h) / ori_h));
     pad_x = (new_w - ori_w * scale) / 2;
     pad_y = (new_h - ori_h * scale) / 2;
     for (unsigned int i = 0; i < out_boxes.size(); i++)
     {
         out_boxes[i].box.x0 = round(std::max((out_boxes[i].box.x0 - pad_x) / scale, float(0)));
         out_boxes[i].box.x1 = round(std::min((out_boxes[i].box.x1 - pad_x) / scale, float(ori_w - 1)));
         out_boxes[i].box.y0 = round(std::max((out_boxes[i].box.y0 - pad_y) / scale, float(0)));
         out_boxes[i].box.y1 = round(std::min((out_boxes[i].box.y1 - pad_y) / scale, float(ori_h - 1)));
     }
 }
 
 void run_model(magik::uranus::Network *model, magik::uranus::Tensor &img, std::vector<magik::uranus::Tensor> &outputs) {
     int img_size = img.get_bytes_size();
     auto input_names = model->get_input_names();
     for(auto n: input_names) {
        //  std::cout<<"input_names: "<< n <<std::endl;
         auto input = model->get_input(n);
         uint8_t *indata = (uint8_t*)(input.vdata<uint8_t>());
         memcpy((void *)indata, (void *)img.vdata<uint8_t>(), img_size * sizeof(uint8_t));
 
         // const unsigned char* in_ptr = input.vdata<uint8_t>();
         // FILE *fp = fopen(std::string(n+".res").c_str(), "wb");
         // if (fp) {
         //     fwrite(in_ptr, 1, (input.get_bytes_size()), fp);
         //     fclose(fp);
         // } else {
         //     printf("open %s file failed\n", std::string(n+".res").c_str());
         //     return;
         // }
     }
 
     // net forword
     int RUN_CNT = 1;
     for (int i = 0; i < RUN_CNT; i++) {
         model->run();
     }
 
     // get outputs
     auto output_names = model->get_output_names();
     for(auto n: output_names) {
        //  std::cout<<"output_names: "<< n <<std::endl;
         auto output = model->get_output(n);
         outputs.push_back(output);
 
         // const unsigned char* out_ptr = output.vdata<uint8_t>();
         // FILE *fp = fopen(std::string(n+".res").c_str(), "wb");
         // if (fp) {
         //     fwrite(out_ptr, 1, (output.get_bytes_size()), fp);
         //     fclose(fp);
         // } else {
         //     printf("open %s file failed\n", std::string(n+".res").c_str());
         //     return;
         // }
     }
 
     // std::string cls_stride_8 = "851_last_layer";
     // std::string cls_stride_16 = "874_last_layer";
     // std::string cls_stride_32 = "897_last_layer";
     // std::string bbox_stride_8 = "840_last_layer";
     // std::string bbox_stride_16 = "863_last_layer";
     // std::string bbox_stride_32 = "886_last_layer";
 
     // auto cls_stride_8_output = model->get_output(cls_stride_8);
     // auto bbox_stride_8_output = model->get_output(bbox_stride_8);
     // auto cls_stride_16_output = model->get_output(cls_stride_16);
     // auto bbox_stride_16_output = model->get_output(bbox_stride_16);
     // auto cls_stride_32_output = model->get_output(cls_stride_32);
     // auto bbox_stride_32_output = model->get_output(bbox_stride_32);
 
     // // Push the outputs to the 'outputs' vector in the order of classification heads followed by detection heads
     // // Note 1: Prioritize classification heads before detection heads
     // // Note 2: The order of elements corresponds to their respective strides
     // outputs.push_back(cls_stride_8_output);
     // outputs.push_back(bbox_stride_8_output);
     // outputs.push_back(cls_stride_16_output);
     // outputs.push_back(bbox_stride_16_output);
     // outputs.push_back(cls_stride_32_output);
     // outputs.push_back(bbox_stride_32_output);
 
     return;
 }
 
 bool is_file(const std::string& path) {
     struct stat buf;
     stat(path.c_str(), &buf);
     return S_ISREG(buf.st_mode);
 }
 
 bool is_directory(const std::string& path) {
     struct stat buf;
     stat(path.c_str(), &buf);
     return S_ISDIR(buf.st_mode);
 }
 
 std::vector<std::string> get_image_files(const std::string& dir_path) {
     std::vector<std::string> image_files;
     DIR* dir = opendir(dir_path.c_str());
     if (!dir) {
         std::cerr << "无法打开目录: " << dir_path << std::endl;
         return image_files;
     }
     
     struct dirent* entry;
     while ((entry = readdir(dir)) != nullptr) {
         std::string filename = entry->d_name;
         if (filename == "." || filename == "..") continue;
         
         std::string path = dir_path + "/" + filename;
         if (is_file(path)) {
             // 检查文件扩展名是否为图片格式
             std::string ext = filename.substr(filename.find_last_of(".") + 1);
             std::transform(ext.begin(), ext.end(), ext.begin(), ::tolower);
             if (ext == "jpg" || ext == "jpeg" || ext == "png" || ext == "bmp" || ext == "gif") {
                 image_files.push_back(path);
             }
         }
     }
     closedir(dir);
     return image_files;
 }
 
 // 修改run_detect_single函数声明，接收模型指针而非模型路径
 bool run_detect_single(magik::uranus::Network* model, std::string img_path, int img_w, int img_h, const std::string& output_dir);

 int main(int argc, char **argv) {
     if (argc != 4) {
         printf("用法: %s <模型路径> <图片文件夹路径> <输出文件夹路径>\n", argv[0]);
         exit(0);
     }
     
     std::string model_path = argv[1];
     std::string img_dir = argv[2];
     std::string output_dir = argv[3];
     
     // 检查输入是否为目录
     if (!is_directory(img_dir)) {
         std::cerr << "错误: " << img_dir << " 不是一个有效的目录" << std::endl;
         return -1;
     }
     
     // 如果输出文件夹存在则删除重建
     struct stat info;
     if (stat(output_dir.c_str(), &info) == 0) {
         // 目录存在，执行删除命令
         std::string rm_cmd = "rm -rf " + output_dir;
         if (system(rm_cmd.c_str()) != 0) {
             std::cerr << "错误: 删除现有输出目录失败: " << output_dir << std::endl;
             return -1;
         }
     }
     // 创建新的输出目录
     if (mkdir(output_dir.c_str(), 0777) != 0) {
         std::cerr << "错误: 创建输出目录失败: " << output_dir << " " << strerror(errno) << std::endl;
         return -1;
     }
     
     // 获取目录中所有图片文件
     std::vector<std::string> image_files = get_image_files(img_dir);
     if (image_files.empty()) {
         std::cerr << "错误: 在目录 " << img_dir << " 中未找到图片文件" << std::endl;
         return -1;
     }
     
     /******************在main中提前加载模型***********************/
     Network **derived_models;
     int model_num = 0;
     printf("load model: %s \n", model_path.c_str());
     if (magik_load_model(derived_models, model_num, model_path.c_str())) {
         printf("create model failed.\n");
         return -1;
     }
     auto model = derived_models[0];
     int ret = model->get_forward_memory_size();
     std::cout << "model forward memory size: " << ret << " = " << ret / 1024 << "KB" << " = " << ret / 1024 / 1024 << "MB" << std::endl;

     // 批量处理图片并统计识别率
     int total_images = image_files.size();
     int detected_images = 0;
     
     std::cout << "开始处理 " << total_images << " 张图片..." << std::endl;
     
     for (const std::string& img_path : image_files) {
        //  std::cout << "\n处理图片: " << img_path << std::endl;
         ori_img_w = -1;
         ori_img_h = -1;
         
         // 传递模型指针给run_detect_single，不再传递模型路径
         if (run_detect_single(model, img_path, ori_img_w, ori_img_h, output_dir)) {
             detected_images++;
             std::cout << "图片 " << img_path << " 检测到目标" << std::endl;
         } else {
             std::cout << "图片 " << img_path << " 未检测到目标" << std::endl;
         }
     }

     /******************所有图片处理完成后销毁模型***********************/
     magik_destroy_model(model);
     printf("magik_destroy_model over\n");
     
     // 计算并输出识别率
     float detection_rate = static_cast<float>(detected_images) / total_images * 100;
     std::cout << "\n处理完成!" << std::endl;
     std::cout << "总图片数: " << total_images << std::endl;
     std::cout << "检测到目标的图片数: " << detected_images << std::endl;
     std::cout << "识别率: " << detection_rate << "%" << std::endl;
     
     return 0;
 }
 
 // 修改run_detect_single函数实现，移除模型加载/销毁逻辑
 bool run_detect_single(magik::uranus::Network* model, std::string img_path, int img_w, int img_h, const std::string& output_dir) {
     /******************移除模型加载代码***********************/
     
     /******************prepare data***********************/
     magik::uranus::shape_t working_img_shape = {1, input_h, input_w, 4};
     magik::uranus::Tensor working_img(working_img_shape);
     printf("load image: %s \n", img_path.c_str());
     FrameInfo input_info;
     Img img;
 
    printf("read_img=========");
    read_img(input_info, img, img_path.c_str());
     
     RET_CHECKER(input_info.data != nullptr, "Read image failed!!!!");
     prepare_data(input_info, working_img);
     /******************model forward***********************/
     std::vector<magik::uranus::Tensor> outputs;
     run_model(model, working_img, outputs);  // 使用外部传入的model

     /******************postprocess***********************/
     std::vector<magik::uranus::ObjBbox_t> temp_boxes;
     vector<magik::uranus::ObjectBbox_t> output_boxes;
     generateBBox(outputs, temp_boxes);
     if (classes == 1) {
         magik::uranus::nms(temp_boxes, output_boxes, nms_threshold, magik::uranus::NmsType_t::HARD_NMS);
     } else {
         manyclass_nms(temp_boxes, output_boxes, classes, 0);
     }
     scale_boxes(input_info.width, input_info.height, input_w, input_h, output_boxes);
     int num_boxes = output_boxes.size();
     std::cout << "num_boxes: " << num_boxes << std::endl;
     
     // 绘制边界框并保存结果
     bool detected = num_boxes > 0;
     for (int i = 0; i < int(output_boxes.size()); i++) {
         auto output_box = output_boxes[i];
         printf("box:   ");
         printf("x0: %d ", (int)output_box.box.x0);
         printf("y0: %d ", (int)output_box.box.y0);
         printf("x1: %d ", (int)output_box.box.x1);
         printf("y1: %d ", (int)output_box.box.y1);
         printf("score: %.2f ", output_box.score);
         printf("class: %d ", output_box.class_id);
         printf("\n");
         Point pt1 = {
             .x = (int)output_box.box.x0,
             .y = (int)output_box.box.y0};
         Point pt2 = {
             .x = (int)output_box.box.x1,
             .y = (int)output_box.box.y1};
         if (DATA_FORMAT == magik::uranus::ChannelLayout::RGBA) {
             sample_draw_box_for_image(&img, pt1, pt2, color, 2);
         }
     }
     
     // 创建输出目录
     mkdir(output_dir.c_str(), 0777);
     
     // 保存结果图片，使用原图片名
     if (DATA_FORMAT == magik::uranus::ChannelLayout::RGBA) {
         // 获取原图片文件名
         size_t last_slash = img_path.find_last_of("/\\");
         std::string img_name = img_path.substr(last_slash + 1);
         std::string result_path = output_dir + "/" + img_name;
         
         // 如果是JPG格式，保留原扩展名
         if (result_path.substr(result_path.find_last_of(".") + 1) != "jpg") {
             result_path = result_path.substr(0, result_path.find_last_of(".")) + ".jpg";
         }
         
         stbi_write_jpg(result_path.c_str(), img.w, img.h, 3, img.data, 75); // w h
         free(img.data);
     }
     
     free(input_info.data);
     /******************移除模型销毁代码***********************/
     
     return detected;
 }
 